#ifndef INDEX_MODEL_TRAINER_H_
#define INDEX_MODEL_TRAINER_H_

#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <netinet/in.h>
#include <netdb.h>  /* netdb is necessary for struct hostent */
#include <fstream>
#include <sstream>
#include <iostream>
#include "config_manager.h"
#include "faiss_index_factory.h"
#include "index_io.h"
#include "json/json.h"

struct ArgvCnt{
    std::string s_program_dir;
    std::string s_train_filename;
    int ui_index_type;
};

struct DataUniqueFlag
{
    int i_appid_;           // application unique identification
    int i_field_id_;        // fieldname unique id

    DataUniqueFlag()
        : i_appid_(-1)
        , i_field_id_(-1)
    { }

    DataUniqueFlag(
            int i_appid,
            int i_field_id)
        : i_appid_(i_appid)
        , i_field_id_(i_field_id)
    { }

    bool operator<(const DataUniqueFlag& o_unique_data) const {
        if (i_appid_ == o_unique_data.i_appid_) {
                return i_field_id_ < o_unique_data.i_field_id_;
        } else {
            return i_appid_ < o_unique_data.i_appid_;
        }
    }
};

int main(int argc, char *argv[])
{
    if (argc != 2 && argc != 3) {
        printf("Usage: %d <train data filename> {<index_type_id>}\n", argc);
        exit(1);
    }

    ArgvCnt o_argv_cnt;
    o_argv_cnt.s_program_dir = argv[0];
    o_argv_cnt.s_train_filename = argv[1];
    o_argv_cnt.ui_index_type = 0;

    if (3 == argc) {
        o_argv_cnt.ui_index_type = atoi(argv[2]);
    }

    bool b_ret = vectorindex::ConfigManager::Instance()->InitConfig();
    if (!b_ret) { return -1; }

    vectorindex::IndexFactory* p_idx_fac = vectorindex::IndexFactory::Instance();
    b_ret = p_idx_fac->InitIndex(true);
    if (!b_ret) { return -1; }

    std::ifstream inf;
    inf.open(o_argv_cnt.s_train_filename);
    if (inf.is_open() == false) {
        printf("open file error: %s.\n", o_argv_cnt.s_train_filename.c_str());
        return -1;
    }

    std::map<DataUniqueFlag , std::vector<float>> o_train_data_map;
    std::string s,sender;
    while (getline(inf, s)) {
        sender = s;

        Json::Reader reader;
        Json::Value send_json_value;
        DataUniqueFlag o_uni_data_idx;
        if (reader.parse(sender , send_json_value)) {
            if (send_json_value.isMember("appId")
                && send_json_value["appId"].isInt()) {
                o_uni_data_idx.i_appid_ = send_json_value["appId"].asUInt();
            } else {
                std::cout << "appId is incorrect" << std::endl;
                getchar();
                return -1;
            }

            if (send_json_value.isMember("fieldId")
                && send_json_value["fieldId"].isInt()) {
                o_uni_data_idx.i_field_id_ = send_json_value["fieldId"].asUInt();
            } else {
                std::cout << "fieldId is incorrect" << std::endl;
                getchar();
                return -1;
            }

            if (send_json_value.isMember("defaultIndexType")
                && send_json_value["defaultIndexType"].isInt()) {
                if (o_argv_cnt.ui_index_type != (int)send_json_value["defaultIndexType"].asUInt()) {
                    continue;
                }
            } else {
                std::cout << "defaultIndexType is incorrect" << std::endl;
                getchar();
                return -1;
            }

            if (send_json_value.isMember("vectorData")
                && send_json_value["vectorData"].isArray()) {
                    Json::Value o_vector_data_array = send_json_value["vectorData"];
                    for (int i = 0; i < (int)o_vector_data_array.size(); i++) {
                        if (o_vector_data_array[i].isDouble()) {
                            o_train_data_map[o_uni_data_idx].emplace_back(o_vector_data_array[i].asDouble());
                        }
                    }
            } else {
                std::cout << "vectorData is incorrect" << std::endl;
                getchar();
                return -1;
            }
        }

        s = "";
    }

    if (o_train_data_map.empty()) {
        std::cout << "train data is empty, index_type_id:" << o_argv_cnt.ui_index_type << std::endl;
        return -1;
    }
    
    std::map<UniqueIndexFlag , faiss::Index*> o_train_idx = p_idx_fac->GetTrainIdxIns();
    
    std::map<DataUniqueFlag , std::vector<float>>::iterator data_iter = o_train_data_map.begin();
    for ( ; data_iter != o_train_data_map.end(); ++data_iter) {
        std::map<UniqueIndexFlag , faiss::Index*>::iterator iter = o_train_idx.begin();
        for ( ; iter != o_train_idx.end(); ++iter) {
            if (data_iter->first.i_appid_ == iter->first.i_appid_ &&
            data_iter->first.i_field_id_ == iter->first.i_field_id_) {
            if (!iter->second->is_trained) {
                std::cout<< "Need train , appid:" << iter->first.i_appid_ << ",fieldid:" << iter->first.i_field_id_ 
                << ",indextypeid:" << iter->first.i_index_typeid_ <<std::endl;
                iter->second->verbose = true;
                std::cout<< data_iter->second.size() <<std::endl;
                iter->second->train(data_iter->second.size() , data_iter->second.data());
                std::string s_save_path = std::to_string(iter->first.i_appid_) + "_" 
                    + std::to_string(iter->first.i_field_id_) + "_" 
                    + std::to_string(iter->first.i_index_typeid_) + ".faissindex";
    
                faiss::write_index(iter->second , s_save_path.c_str());
            } else {
                std::cout<< "No Need train , appid:" << iter->first.i_appid_ << ",fieldid:" << iter->first.i_field_id_ 
                << ",indextypeid:" << iter->first.i_index_typeid_ <<std::endl;
            }
            }
        }
    }
};

#endif
